package fun

import (
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"net/http"
	"strings"
	"sync"
	"ttbot_vx/shared/utils"
)

func init() {
	ModuleList = append(ModuleList, Module{
		func(msg *message) bool {
			if strings.HasPrefix(msg.Content, "/chat") || (msg.IsAt() && strings.HasPrefix(msg.Content, "@"+msg.Owner().NickName)) {
				return true
			}
			return false
		},
		ernieChat,
	})
}

func ernieChat(ctx *messageContext) {
	msg := ctx.Message
	go func() {
		uid := msg.FromUserName
		query := strings.TrimPrefix(strings.TrimPrefix(msg.Content, "/chat"), "@"+msg.Owner().NickName)
		for k, v := range map[string]string{
			"\n": "\\n",
			"\t": "\\t",
			"\r": "\\r",
			"\"": "\\\"",
			"\\": "\\\\",
		} {
			query = strings.ReplaceAll(query, k, v)
		}

		mutex.RLock()
		msgList := msgMap[uid]
		mutex.RUnlock()

		msgList = append(msgList, `{"role":"user","content":"`+query+`"}`)
		resp, err := getResp(`{"messages":[` + strings.Join(msgList, ",") + `]}`)
		if err != nil {
			ctx.ReplyText("出现了一点错误：\n" + err.Error())
			return
		}

		ctx.ReplyText(resp)
		for k, v := range map[string]string{
			"\n": "\\n",
			"\t": "\\t",
			"\r": "\\r",
			"\"": "\\\"",
			"\\": "\\\\",
		} {
			resp = strings.ReplaceAll(resp, k, v)
		}
		msgList = append(msgList, `{"role":"assistant","content":"`+resp+`"}`)
		mutex.Lock()
		msgMap[uid] = msgList
		mutex.Unlock()
	}()
}

var msgMap = map[string][]string{}
var mutex sync.RWMutex

func getResp(msgListStr string) (string, error) {
	// println("\n" + msgListStr + "\n")

	accessToken, err := getAccessToken(APIKey, SecretKey)
	if err != nil {
		return "", err
	}

	url := "https://aip.baidubce.com/rpc/2.0/ai_custom/v1/wenxinworkshop/chat/eb-instant?access_token=" + accessToken
	// payload := strings.NewReader(`{"messages":[{"role":"user","content":"请你介绍一下你自己"}]}`)
	payload := strings.NewReader(msgListStr)
	client := &http.Client{}
	req, err := http.NewRequest("POST", url, payload)
	if err != nil {
		return "", err
	}

	req.Header.Add("Content-Type", "application/json")
	res, err := client.Do(req)
	if err != nil {
		return "", err
	}

	defer res.Body.Close()
	body, err := io.ReadAll(res.Body)
	if err != nil {
		return "", err
	}
	fmt.Println(string(body))

	respObj := map[string]string{}
	json.Unmarshal(body, &respObj)

	if strings.Contains(string(body), "error_msg") {
		return "", errors.New(respObj["error_msg"])
	} else {
		return respObj["result"], nil
	}
}

func getAccessToken(APIKey, SecretKey string) (string, error) {
	url := "https://aip.baidubce.com/oauth/2.0/token"
	postData := fmt.Sprintf("grant_type=client_credentials&client_id=%s&client_secret=%s", APIKey, SecretKey)
	resp, err := http.Post(url, "application/x-www-form-urlencoded", strings.NewReader(postData))
	if err != nil {
		return "", err
	}

	defer resp.Body.Close()
	body, err := io.ReadAll(resp.Body)
	if err != nil {
		return "", err
	}

	accessTokenObj := map[string]string{}
	json.Unmarshal([]byte(body), &accessTokenObj)
	return accessTokenObj["access_token"], nil
}

var APIKey string = utils.Config.ErineBot.APIKey
var SecretKey string = utils.Config.ErineBot.SecretKey
